from PyQt5 import QtCore,QtWidgets
from configparser import ConfigParser
import gym_env
import gym
import os
from stable_baselines3 import SAC
import numpy as np
import sys
from ui_train import TrainingUi

class EvaluateThread(QtCore.QThread):
    def __init__(self, config_path):
        super(EvaluateThread, self).__init__()
        self.cfg = ConfigParser()
        self.cfg.read(config_path)

        self.env = gym.make('airsim-env-v0',cfg=self.cfg)
        self.env.set_config(self.cfg)
        self.project_path = os.path.dirname(config_path)

    def run(self):
        model = SAC.load(f"{self.project_path}/models/model.zip", env=self.env)
        self.env.model = model
        obs = self.env.reset()
        episode_num = 0
        time_step = 0
        reward_sum = np.array([.0])
        episode_successes = []
        episode_crashes = []

        traj_list = []
        action_list = []
        state_raw_list = []
        step_num_list = []
        obs_list = []
        while True:
            unscaled_action, _ = model.predict(obs, deterministic=True)
            time_step +=1
            new_obs,reward,done,info = self.env.step(unscaled_action)
            pose = self.env.dynamic_model.get_position()
            traj_list.append(pose)
            action_list.append(unscaled_action)
            state_raw_list.append(self.env.dynamic_model.state_raw)
            obs_list.append(obs)
            
            obs = new_obs
            reward_sum[-1] += reward

            if done:
                episode_num += 1
                maybe_is_success = info.get('is_success')
                maybe_is_crash = info.get('is_crash')
                print('episode: ', episode_num, ' reward:', reward_sum[-1],
                      'success:', maybe_is_success)
                episode_successes.append(float(maybe_is_success))
                episode_crashes.append(float(maybe_is_crash))
                reward_sum = np.append(reward_sum, .0)
                obs = self.env.reset()
                if info.get('is_success'):
                    traj_list.append(1)
                    action_list.append(1)
                    step_num_list.append(info.get('step_num'))
                elif info.get('is_crash'):
                    traj_list.append(2)
                    action_list.append(2)
                else:
                    traj_list.append(3)
                    action_list.append(3)
                traj_list = []
                action_list = []
                state_raw_list = []
                obs_list = []

def main():
    config_path = "result/NewYork_2024_05_18_12_03/config.ini"
    app = QtWidgets.QApplication(sys.argv)
    gui = TrainingUi(config=config_path)
    gui.show()

    evaluate_thread = EvaluateThread(config_path)
    evaluate_thread.env.action_signal.connect(gui.action_cb)
    evaluate_thread.env.state_signal.connect(gui.state_cb)
    evaluate_thread.env.attitude_signal.connect(gui.attitude_plot_cb)
    evaluate_thread.env.reward_signal.connect(gui.reward_plot_cb)
    evaluate_thread.env.pose_signal.connect(gui.traj_plot_cb)

    evaluate_thread.start()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()